#include "ML_DecisionStumpPreprocess.h"
#include "ML_DecisionStumpParams.h"
#include "../Common/ErrorCode.h"
#include "../Math/MathTools.h"
#include "../Tools/ZLogManager.h"

GLOBAL_NAMESPACE_BEGIN
NAMESPACE_MACHINE_LEARNING_BEGIN

MLDecisionStumpPreprocess::MLDecisionStumpPreprocess(int stepNum)
    : stepNum_(stepNum)
    , maxClass_(0)
    , bUseUserDefinedThreshold_(false)
    , nDim_(0)
{

}

MLDecisionStumpPreprocess::~MLDecisionStumpPreprocess()
{

}

void MLDecisionStumpPreprocess::reset()
{
    oneClassStumps_.clear();
    maxCoeffs_.clear();
    minCoeffs_.clear();
    thresholdCandidate_.clear();
    bUseUserDefinedThreshold_ = false;
    maxClass_ = 0;
    stepNum_ = 10;
    nDim_ = 0;
}

void MLDecisionStumpPreprocess::setStepNum(int stepNum)
{
    stepNum_ = stepNum;
    bUseUserDefinedThreshold_ = false;
}

void MLDecisionStumpPreprocess::setThresholdCandidate(const std::vector<double>& thresholdCandidates)
{
    if (thresholdCandidates.size() <= 1)
        return;
    thresholdCandidate_ = thresholdCandidates;
    stepNum_ = thresholdCandidates.size();
    bUseUserDefinedThreshold_ = true;
}
void MLDecisionStumpPreprocess::setMaxClassNum(int maxClass)
{
    maxClass_ = maxClass;
}

int MLDecisionStumpPreprocess::getClassSize()
{
    return maxClass_;
}

int MLDecisionStumpPreprocess::getDimensionSize()
{
    return nDim_;
}

int MLDecisionStumpPreprocess::getThresholdSize()
{
    return stepNum_;
}


double MLDecisionStumpPreprocess::getThreshold(int idx, int dim)
{
    if (dim == -1)
    {
        return thresholdCandidate_[idx];
    }
    else
        return minCoeffs_[dim] + (maxCoeffs_[dim] - minCoeffs_[dim]) / (getThresholdSize()-1)*idx;
}

void MLDecisionStumpPreprocess::generateThresholds(const Eigen::MatrixXd& trainData)
{
    size_t nDataSize = trainData.rows();
    size_t nDataDim = getDimensionSize();
    maxCoeffs_.resize(nDataDim, 0);
    minCoeffs_.resize(nDataDim, 0);
    for (size_t i = 0; i < nDataDim; i++)
    {
        maxCoeffs_[i] = trainData.col(i).maxCoeff();
        minCoeffs_[i] = trainData.col(i).minCoeff();
    }
    bUseUserDefinedThreshold_ = false;
}

void MLDecisionStumpPreprocess::generateParams()
{
    for (int c = 0; c < getClassSize(); c++)
    {
        for (int d = 0; d < getDimensionSize(); d++)
        {
            for (int i = 0; i < getThresholdSize(); i++)
            {
                MLDecisionStumpParams* stump = getStump(c, d, i, true);
                stump->subClass.clear();
                stump->subClass.push_back(c);
                stump->dimF = d;
                stump->threshold = getThreshold(i, bUseUserDefinedThreshold_ ? -1 : d);
                stump->threIdx = i;
                stump->k.resize(maxClass_, 0);
                stump->bGreater = true;

                stump = getStump(c, d, i, false);
                stump->subClass.clear();
                stump->subClass.push_back(c);
                stump->dimF = d;
                stump->threshold = getThreshold(i, bUseUserDefinedThreshold_ ? -1 : d);
                stump->threIdx = i;
                stump->k.resize(maxClass_, 0);
                stump->bGreater = false;
            }
        }
    }
}

int MLDecisionStumpPreprocess::preProcess(const Eigen::MatrixXd& trainData, const Eigen::MatrixXi& labels, const Eigen::MatrixXd& weights)
{
    int ret = ErrorCode::ERROR_NO_Error;

    if (getClassSize() <= 0)
    {
        DebugLog << "Make sure to be defined the maxClassIdx first.";
        return ErrorCode::ERROR_Invalid_Parameter;
    }

    // threshold
    nDim_ = trainData.cols();

    oneClassStumps_.resize(getClassSize()*getDimensionSize()*getThresholdSize()*2);
    if (!bUseUserDefinedThreshold_)
        generateThresholds(trainData);
    generateParams();

    for (size_t i = 0; i < oneClassStumps_.size(); i++)
    {
        trainOneStump(&oneClassStumps_[i], trainData, labels, weights);
    }

    return ret;
}

MLDecisionStumpParams* MLDecisionStumpPreprocess::getStump(int classC, int dimF, int threIdx, bool greater)
{
    int idx = (classC*(getDimensionSize()*getThresholdSize()) + dimF*getThresholdSize() + threIdx)*2+(greater ? 0 : 1);
    return &oneClassStumps_[idx];
}

int MLDecisionStumpPreprocess::getLabelResponse(MLDecisionStumpParams* stump, const Eigen::RowVectorXi& label)
{
    // assume the stump is a simple stump (only has one class)
    int classC = *stump->subClass.begin();
    return label(classC) == 1 ? 1 : -1;
}

int MLDecisionStumpPreprocess::trainOneStump(MLDecisionStumpParams* stump, const Eigen::MatrixXd& trainData, const Eigen::MatrixXi& labels, const Eigen::MatrixXd& weights)
{
    int ret = ErrorCode::ERROR_NO_Error;
    int nDataSize = trainData.rows();
    int nDataDim = trainData.cols();

    double numeratorP = 0, denominatorP = 0;
    double numeratorN = 0, denominatorN = 0;
    for (size_t i = 0; i < nDataSize; i++)
    {
        int predict = getLabelResponse(stump, labels.row(i));
        double value = trainData(i, stump->dimF);
        double weight = weights(i, stump->dimF);
        if ((stump->bGreater && value > stump->threshold)
            || (!stump->bGreater && value < stump->threshold))
        {
            numeratorP   += weight * predict;
            denominatorP += weight;
        }
        else
        {
            numeratorN   += weight * predict;
            denominatorN += weight;
        }
    }
    stump->a = MathTools::isZero(denominatorP) ? 0 : numeratorP / denominatorP;
    stump->b = MathTools::isZero(denominatorN) ? 0 : numeratorN / denominatorN;
    stump->wcP = denominatorP;
    stump->wcN = denominatorN;

    for (int c = 0; c < getClassSize(); c++)
    {
        double numeratorQ = 0, denominatorQ = 0;
        for (size_t i = 0; i < nDataSize; i++)
        {
            int predict = labels(i, c) == 1 ? 1 : -1;
            double weight = weights(i, c);
            numeratorQ += weight * predict;
            denominatorQ += weight;
        }
        stump->k[c] = MathTools::isZero(denominatorQ) ? 0 : numeratorQ / denominatorQ;
    }

    return ret;
}


NAMESPACE_MACHINE_LEARNING_END
GLOBAL_NAMESPACE_END
